#!/bin/env python2
# -*- coding: utf-8 -*-
import os, sys
from utils import read_articles, write_article, uopen, AttrDict, filenames, open_corpus
from multiprocessing import Pool
from itertools import izip
from functools import partial
from os import path
from alignment import AlignmentModel
from munteanu import MunteanuModel
from random import shuffle

def munteanu_process(corpus, suffix, giza_model, training_corpus, params):
    input_corpus = corpus
    output_corpus = '{}.{}'.format(corpus, suffix)
    params = AttrDict(params.items())

    alignment_model = AlignmentModel(giza_model, params)
    print '{}: model initialized'.format(corpus)
    munt_model = MunteanuModel(alignment_model, training_corpus, params)
    print '{}: model trained'.format(corpus)

    src_filename, trg_filename = filenames(input_corpus, params)
    src_articles = read_articles(src_filename)
    trg_articles = read_articles(trg_filename)

    with open_corpus(output_corpus, params, 'w') as (src_output_file, trg_output_file), \
         open(output_corpus + '.scores', 'w') as score_output_file:
        print '{}: files created'.format(corpus)
        
        for (index, src_txt), (_, trg_txt) in izip(src_articles, trg_articles):
            pairs = munt_model.extract_pairs(src_txt.split('\n'), trg_txt.split('\n'))
            if not pairs:
                continue
            src_sentences, trg_sentences, scores = zip(*pairs)

            # MaxEnt scores are saved in a .scores file
            score_article = '\n'.join('{:.4f}'.format(score) for score in scores) + '\n'
            src_article = '\n'.join(src_sentences) + '\n'
            trg_article = '\n'.join(trg_sentences) + '\n'

            write_article(score_article, index, score_output_file)
            write_article(src_article, index, src_output_file)
            write_article(trg_article, index, trg_output_file)

    print '{}: finished'.format(corpus)

def main():
    help_msg = """
    Runs MUNT in parallel on several document pairs. The output corpus is named like the input corpus, with a suffix.

    Usage:
        {prog} SUFFIX GIZA_MODEL TRAINING_CORPUS SRC_LANG TRG_LANG FILES
    Example of usage:
        {prog} munt ~/models/europarl ~/data/news.train fr en ~/output/*.en ~/output/*.fr

    SUFFIX: suffix to append to the name of the input corpora.
    GIZA_MODEL: path of the GIZA++ model to use in MUNT.
    TRAINING CORPUS: path of the parallel corpus used to train MUNT.
    SRC_LANG, TRG_LANG: file extensions for the source and target sides of the corpus.
    FILES: list of file names on which to apply MUNT (source and target documents).
    """
    try:
        suffix, giza_model, training_corpus, src_lang, trg_lang = sys.argv[1:6]
        files = sys.argv[6:]
    except:
        sys.exit(help_msg.format(prog=sys.argv[0]))

    # Settings:
    params = AttrDict(n_positive=800, n_negative=2400, threshold=0.95, keep_score=True, src_lang=src_lang, trg_lang=trg_lang)

    partial_munt = partial(munteanu_process, suffix=suffix, giza_model=giza_model,
            training_corpus=training_corpus, params=params)

    src_files = set([f.rsplit('.', 1)[0] for f in files if f.endswith('.' + src_lang)])
    trg_files = set([f.rsplit('.', 1)[0] for f in files if f.endswith('.' + trg_lang)])

    files = src_files.intersection(trg_files)

    p = Pool()
    p.map(partial_munt, files)

if __name__ == '__main__':
    main()

